

#include "membackend/accurateController.h"

Controller::Controller() : transaction_queue(), cmd_queue()
{
    m_numRanks = 0;
    m_numBanks = 0;
    m_numBankGroups = 0;

    trans_id = 0;
    clock_cycle = 0;
    is_draining = 0; // Initialize draining flag
}

Controller::Controller(AddressMappingConfig *amp) : transaction_queue(amp), cmd_queue(amp)
{

    m_addressMappingConfig = amp;
    m_numRanks = amp->getRanks();
    m_numBanks = amp->getBanksPerGroup();
    m_numBankGroups = amp->getBankGroups();

    this->trans_queue_size = amp->getTransQueueSize();

    trans_id = 0;
    clock_cycle = 0;
    is_draining = 0; // Initialize draining flag

    std::cout << "Controller created with " << m_numRanks << " ranks, "
              << m_numBanks << " banks per group, and "
              << m_numBankGroups << " bank groups." << std::endl;

    m_ranksVector.reserve(m_numRanks);
    for (uint64_t i = 0; i < m_numRanks; ++i)
    {
        m_ranksVector.emplace_back(m_numBanks, m_numBankGroups);
    }
}

int Controller::Init()
{

    read_transactions.clear();   // Clear previous read transactions
    write_transactions.clear();  // Clear previous write transactions
    return_transactions.clear(); // Clear previous return transactions

    read_transactions.reserve(trans_queue_size);
    write_transactions.reserve(trans_queue_size);

    // std::cout << "Initializing Controller with " << std::endl;
    // Initialization logic for each rank
    for (uint64_t i = 0; i < m_numRanks; ++i)
    {
        if (m_ranksVector[i].Init(this->m_addressMappingConfig) != 0)
        {
            return -1; // Initialization failed
        }
    }

    return 0;
}

// 读的冲突

void Controller::Clock(uint64_t cycle)
{
    clock_cycle = cycle;
    std::cout << std::endl
              << "------------------------------------------" << std::endl;
    std::cout << "Clock cycle: " << clock_cycle << std::endl;
}

bool Controller::WillAcceptTransaction(bool isWrite)
{
    bool result = true;
    if (isWrite)
    {
        result = write_transactions.size() < write_transactions.capacity(); // Example limit for write transactions
        std::cout << "WillAcceptTransaction: Write transactions size: " << write_transactions.size() << ", capacity: " << write_transactions.capacity() << ", result: " << result << std::endl;
    }
    else
    {
        result = read_transactions.size() < read_transactions.capacity(); // Example limit for read transactions
        std::cout << "WillAcceptTransaction: Read transactions size: " << read_transactions.size() << ", capacity: " << read_transactions.capacity() << ", result: " << result << std::endl;
    }
    return result;
}

bool Controller::AddTransaction(uint64_t reqid, uint64_t address, bool isWrite)
{

    bool result = false;

    TRANSACTION_TYPE type = isWrite ? TRANSACTION_TYPE::WRITE : TRANSACTION_TYPE::READ;
    Transaction transaction(reqid, address, type, this->getAddressMappingConfig());
    transaction.setStartTime(this->clock_cycle);
    transaction.setEndTime(this->clock_cycle + 1);

    if (type == TRANSACTION_TYPE::READ)
    {

        if (pending_writes.find(address) != pending_writes.end())
        {
            return_transactions.push_back(transaction);
        }
        else
        {

            if (pending_reads.find(address) != pending_reads.end())
            {
                pending_reads[address].push_back(transaction);
            }
            else
            {
                vector<Transaction> pending_read_vector;
                pending_read_vector.push_back(transaction);
                pending_reads[address] = pending_read_vector;
            }

            all_transactions[reqid] = transaction;
            read_transactions.push_back(transaction);
        }

        result = true;
    }
    else
    {

        result = false;
    }

    return result;
}

Rank& Controller::getRank(uint64_t rank_id)
{
    if (rank_id < m_numRanks)
    {
        return m_ranksVector[rank_id];
    }
    else
    {
        throw std::out_of_range("Rank ID out of range");
    }
}

Bank& Controller::getBank(uint64_t rank_id, uint64_t bank_group_id, uint64_t bank_id)
{
    std::cout << "Controller::getBank called with rank_id: " << rank_id 
              << ", bank_group_id: " << bank_group_id 
              << ", bank_id: " << bank_id << std::endl;
    return this->getRank(rank_id).getBank(bank_group_id, bank_id);
}

/*
一行调一个函数


*/

bool Controller::isNeedDrain()
{
    bool result = false;
    if ((write_transactions.size() >= write_transactions.capacity()) || (write_transactions.size() > 0 && cmd_queue.isEmpty()))
    {
        result = true; // 需要排空
    }

    return result;
}



uint64_t Controller::Process(uint64_t addr, bool isWrite, uint64_t reqID, uint64_t time)
{
    uint64_t delay = 0;
    bool flag = 0;

    // std::cout << "isWrite: " << isWrite << std::endl;

    TRANSACTION_TYPE type = isWrite ? TRANSACTION_TYPE::WRITE : TRANSACTION_TYPE::READ;
    std::cout << "AccurateDRAM issueRequest type: " << (type == TRANSACTION_TYPE::READ ? "READ" : "WRITE") << std::endl;

    Transaction transaction(reqID, addr, type, this->getAddressMappingConfig());
    std::cout << "Controller::Process: Processing transaction with ID: " << reqID << " at address: " << std::hex << addr << std::dec << ", rank_id: " << transaction.getRankId() 
              << ", bank_group_id: " << transaction.getBankGroupId() 
              << ", bank_id: " << transaction.getBankId() 
              << ", row_id: " << transaction.getRowId()
              << ", column_id: " << transaction.getColumnId()
              << std::endl;

    transaction.setStartTime(time);
    transaction.setEndTime(time + 1);

    if (type == TRANSACTION_TYPE::READ)
    {
        // 写后读命中，模拟 bypass，直接返回
        if (pending_writes.find(addr) != pending_writes.end())
        {
            transaction.setLatency(1);
            transaction.setEndTime(time + 1);
            return_transactions.push_back(transaction);
            delay += 1;

            std::cout << "Controller::Process: Read transaction bypassed write buffer with ID: " << reqID << " at time: " << time << " and delay is: " << std::dec << delay << std::endl;
            return delay;
        }
        else
        {
            if (pending_reads.find(addr) != pending_reads.end())
            {
                flag = true;
                vector<Transaction> &pending_read_vector = pending_reads[addr];
                
                std::cout << "Pending read transactions for address: " << std::hex << addr << std::dec << "--------------------------------" << std::endl;
                for(auto &pending_read : pending_read_vector)
                {
                    std::cout << "Pending read transaction ID: " << pending_read.getId() << " at address: " << std::hex << addr << std::dec << std::endl;
                }

                std::cout << "------------------------------------------" << std::endl;
                uint64_t delay = pending_read_vector.back().getLatency();
                std::cout << "delay: " << delay << std::endl;
                transaction.setLatency(delay);
                transaction.setEndTime(time + delay);
                pending_reads[addr].push_back(transaction);
                return_transactions.push_back(transaction);
                all_transactions[reqID] = transaction;
                std::cout << "发生读后读命中，地址为: " << std::hex << addr << std::dec << " delay is: " << delay << std::endl;

                return delay+1;
            }
            else
            {
                flag = false;
                vector<Transaction> pending_read_vector;
                pending_read_vector.push_back(transaction);
                pending_reads[addr] = pending_read_vector;
            }

            all_transactions[reqID] = transaction;
            read_transactions.push_back(transaction);
        }
    }
    else
    {

        // 延迟写策略，直接往写缓冲加就好
        all_transactions[reqID] = transaction;
        write_transactions.push_back(transaction);
        pending_writes[addr].push_back(transaction);
        return_transactions.push_back(transaction);

        delay += 1;

        // 这部分应该放在写命令添加之后判断，还需要考虑是否向命令队列添加激活和刷新指令
        if (this->isNeedDrain())
        {
            is_draining = write_transactions.size();
            uint64_t temp_delay = 0;
            uint64_t write_delay = 0;
            for (int i = 0; i < is_draining; i++)
            {
                Transaction write_transaction = write_transactions[i];
                std::cout << "Processing write transaction ID: " << write_transaction.getId() 
                          << " at address: " << std::hex << write_transaction.getAddress() 
                          << std::dec << ", rank_id: " << write_transaction.getRankId() 
                          << ", bank_group_id: " << write_transaction.getBankGroupId() 
                          << ", bank_id: " << write_transaction.getBankId() 
                          << ", row_id: " << write_transaction.getRowId() 
                          << ", column_id: " << write_transaction.getColumnId() 
                          << std::endl;

                std::cout << "**********write_transaction.getId(): " << write_transaction.getId() << std::endl;


                DRAMCommand cmd = cmd_queue.Trans2CMD(write_transaction);
                std::cout << "Processing cmd: " << cmd.getTransId() 
                          << ", type: " << (cmd.type == CommandType::READ ? "READ" : "WRITE") 
                          << ", address: " << std::hex << cmd.address 
                          << std::dec << ", rank_id: " << cmd.rank_id 
                          << ", bank_group_id: " << cmd.bank_group_id 
                          << ", bank_id: " << cmd.bank_id 
                          << ", row_id: " << cmd.row_id 
                          << ", column_id: " << cmd.column_id 
                          << std::endl;
                if (cmd_queue.WillAcceptCommand(cmd.rank_id, cmd.bank_group_id, cmd.bank_id))
                {
                    // cmd_queue.AddCommand(cmd);
                    temp_delay = cmd_queue.Process(cmd, time + write_delay);
                }
                else
                {
                    std::cout << "Command queue is full, cannot accept command for transaction ID: " << write_transaction.getId() << std::endl;
                }

                std::cout << "1 " << std::endl;
                temp_delay += this->getBank(cmd.rank_id, cmd.bank_group_id, cmd.bank_id).Process(cmd, time + temp_delay, this->getAddressMappingConfig());
                write_delay += 1;
            }
        }

        std::cout << "Controller::Process: Write transaction added with ID: " << reqID << " at time: " << std::dec << time << " and delay is: " << delay << std::endl;
        return delay;
    }

    // 如果是读后读命中，直接返回
    if (flag == false)
    {

        delay += this->transaction_queue.Process(transaction, time);
        // 开始处理命令队列

        DRAMCommand cmd = cmd_queue.Trans2CMD(transaction);

        // 改用cmd作为参数进行处理，累加插入到命令队列的延时
        delay += cmd_queue.Process(cmd, time + delay);
        std::cout << "delay after cmd_queue.Process: " << delay << std::endl;
        // 改用cmd作为参数进行处理，累加执行该命令的延时，并更新bank状态；调用getBank方法获取对应的Bank对象

        std::cout << "Processing command for transaction ID: " << "rank_id: " << cmd.rank_id 
                  << ", bank_group_id: " << cmd.bank_group_id 
                  << ", bank_id: " << cmd.bank_id 
                  << ", row_id: " << cmd.row_id 
                  << ", column_id: " << cmd.column_id 
                  << ", type: " << (cmd.type == CommandType::READ ? "READ" : "WRITE") 
                  << ", address: " << std::hex << cmd.address << std::dec << std::endl;
        delay += this->getBank(cmd.rank_id, cmd.bank_group_id, cmd.bank_id).Process(cmd, time + delay, this->getAddressMappingConfig());

        std::cout << "Controller::Process:" << ((cmd.type == CommandType::READ) ? "READ" : "WRITE") << " command processed for transaction ID: " << transaction.getId() << " at time: " << std::dec << time << " with delay: " << delay << std::endl;

        transaction.setLatency(delay);
        transaction.setEndTime(time + delay);
        all_transactions[transaction.getId()].setLatency(delay);
        all_transactions[transaction.getId()].setEndTime(time + delay);
        if (type == TRANSACTION_TYPE::READ)
        {
            read_transactions.back().setLatency(delay);
            read_transactions.back().setEndTime(time + delay);
            pending_reads[addr].back().setLatency(delay);
            pending_reads[addr].back().setEndTime(time + delay);
        }
        else
        {
            write_transactions.back().setLatency(delay);
            write_transactions.back().setEndTime(time + delay);
            pending_writes[addr].back().setLatency(delay);
            pending_writes[addr].back().setEndTime(time + delay);
        }

        return_transactions.push_back(transaction);
    }

    std::cout << "finally delay: " << delay << std::endl;
    return delay;
}

Transaction Controller::getTransactionFromId(uint64_t id)
{
    if (all_transactions.find(id) != all_transactions.end())
    {
        return all_transactions[id];
    }
    else
    {
        throw std::out_of_range("Transaction ID not found");
    }
}

void Controller::DeleteTransactionInAll(uint64_t id)
{
    all_transactions.erase(id);
}

void Controller::DeleteTransactionBuffer(uint64_t id, TransactionType type)
{
    if (type == TransactionType::READ)
    {
        read_transactions.erase(
            std::remove_if(read_transactions.begin(), read_transactions.end(),
                           [id](const Transaction &t)
                           { return t.getId() == id; }),
            read_transactions.end());
    }
    else
    {
        write_transactions.erase(
            std::remove_if(write_transactions.begin(), write_transactions.end(),
                           [id](const Transaction &t)
                           { return t.getId() == id; }),
            write_transactions.end());
    }
}

void Controller::DeletePending(uint64_t id, TransactionType type, uint64_t addr)
{
    if (type == TransactionType::READ)
    {
        if (pending_reads.find(addr) != pending_reads.end())
        {
            auto &pending_read_vector = pending_reads[addr];

            std::cout << "Before delete-----------------------------" << std::endl;
            for (auto i : pending_read_vector)
            {
                std::cout << "Pending read transaction ID: " << i.getId() << " at address: " << std::hex << addr << std::dec << std::endl;
            }
            std::cout << "------------------------------------------" << std::endl;

            pending_read_vector.erase(
                std::remove_if(pending_read_vector.begin(), pending_read_vector.end(),
                               [id](const Transaction &t)
                               { return t.getId() == id; }),
                pending_read_vector.end());

            if (pending_read_vector.empty())
            {
                pending_reads.erase(addr);
                std::cout << "Address " << std::hex << addr << std::dec << " fully removed from pending_reads." << std::endl;
            }

            std::cout << "After delete------------------------------" << std::endl;
            for (auto i : pending_read_vector)
            {
                std::cout << "Pending read transaction ID: " << i.getId() << " at address: " << std::hex << addr << std::dec << std::endl;
            }
            std::cout << "------------------------------------------" << std::endl;
        }
    }
    else
    {
        if (pending_writes.find(addr) != pending_writes.end())
        {
            auto &pending_write_vector = pending_writes[addr];
            pending_write_vector.erase(
                std::remove_if(pending_write_vector.begin(), pending_write_vector.end(),
                               [id](const Transaction &t)
                               { return t.getId() == id; }),
                pending_write_vector.end());

            if (pending_write_vector.empty())
            {
                pending_writes.erase(addr);
                std::cout << "Address " << std::hex << addr << std::dec << " fully removed from pending_writes." << std::endl;
            }
        }
    }
}


void Controller::DeleleReturn(uint64_t id)
{
    return_transactions.erase(
        std::remove_if(return_transactions.begin(), return_transactions.end(),
                       [id](const Transaction &t)
                       { return t.getId() == id; }),
        return_transactions.end());
}

void Controller::DeleteCommandFromQueue(uint64_t rank_id, uint64_t bank_group_id, uint64_t bank_id, uint64_t id)
{
    cmd_queue.DeleteCommand(rank_id, bank_group_id, bank_id, id);
}
